from growthAnalyzer.service.models import TripleResNet50
from collections import OrderedDict
import os
from PIL import Image
import torchvision.transforms as transforms
import torch
from joblib import load

# 加载先前保存的scaler对象
scaler = load('growthAnalyzer/service/scaler.save')


def load_and_transform_image(image_path, convert_to='RGB'):
    """根据提供的路径加载单个图像，并应用转换。

    Args:
        image_path (str): 图像文件的路径。
        convert_to (str): 转换图像的模式，默认为'RGB'。

    Returns:
        torch.Tensor: 转换后的图像张量。
    """
    if not os.path.exists(image_path):
        raise FileNotFoundError(f"The image at {image_path} does not exist.")

    transform = transforms.Compose([
        # 根据需要添加其他转换，例如Resize
        transforms.ToTensor()  # 将PIL图像转换为Tensor
    ])

    with Image.open(image_path) as img:
        if convert_to:
            img = img.convert(convert_to)
        img_tensor = transform(img)

    return img_tensor


# image1 = load_and_transform_image('165_1_TopViewCamera.png')
# print(image1)
# image2 = load_and_transform_image('165_1_SideViewCamera.png')
# image3 = load_and_transform_image('165_1_DepthCamera.png')


# 使用示例
# img_path = 'path_to_your_image.png'
# img_tensor = load_and_transform_image(img_path)


def predict_images(image1, image2, image3):
    model = TripleResNet50()
    # 加载模型权重

    state_dict = torch.load('growthAnalyzer/service/best_MAPE_model.pth', map_location=torch.device('cpu'))

    new_state_dict = OrderedDict()
    for k, v in state_dict.items():
        name = k[7:] if k.startswith('module.') else k
        new_state_dict[name] = v
    model.load_state_dict(new_state_dict)

    # 模型和scaler加载

    scaler = load('growthAnalyzer/service/scaler.save')
    model.eval()

    # 确保输入图像在CPU上并添加批次维度
    image1, image2, image3 = image1.unsqueeze(0), image2.unsqueeze(0), image3.unsqueeze(0)

    # 进行模型预测
    with torch.no_grad():
        outputs = model(image1, image2, image3)

    # 从模型输出获取numpy数组
    predictions = outputs.numpy()

    # 反标准化操作
    predictions_original = scaler.inverse_transform(predictions)

    # 由于outputs可能包含多个预测，确保返回一个包含8个元素的列表
    if predictions_original.size == 8:
        return predictions_original.tolist()
    else:
        raise ValueError("预测值数量不符，确保模型输出8个预测值")

# 使用示例
# 假设image1, image2, image3是处理过的张量，可以这样调用函数：
# result = predict_images(image1, image2, image3)
# print(result)
